from prompts.problem_classification_prompt import prompt_problem_classification, prompt_content_classification
from utils.functions import Get_List, innovation_dict, motherboard_dict, entrepreneurship_dict
from utils.generate_token import api_key
from langchain_core.output_parsers import StrOutputParser
from langchain_community.chat_models import ChatZhipuAI
from langchain_core.prompts import ChatPromptTemplate
from utils.db import LawDB
from collections import OrderedDict
from pprint import pprint
import pymysql
import json
import time
import ast

class Problem_Classification:

    def __init__(self):
        self.llm = ChatZhipuAI(
            api_key=api_key,
            model_name="glm-4-plus",
            streaming=True,
            temperature=0.01,
            top_p=0.01,
        )


    async def classification_problem(self, start_index: int, end_index: int):
        '''
        :param start_index: 开始索引
        :param end_index: 结束索引
        :return:
        '''

        db_laws = pymysql.Connect(
            host='nj-cdb-1mwvwkzt.sql.tencentcdb.com',
            port=63972,
            user='root',
            passwd='grow1234',
            db='stock_exchange'
        )

        test_dict = OrderedDict()
        for problem_id in range(start_index, end_index):
            test_dict[problem_id] = OrderedDict()
            print('正在获取详情页:', problem_id)
            s_time = time.time()

            try:
                cursor = db_laws.cursor()
                cursor.execute("SELECT problem_title, problem_classification, project_id FROM problem_new WHERE problem_id = %s", (problem_id,))
                problem_results = cursor.fetchone()
                问题题目 = problem_results[0]
                问题分类 = problem_results[1]
                # print(1111, len(str(问题分类)))
                if len(str(问题分类)) > 2:
                    continue

                project_id = problem_results[2]
                cursor.execute("SELECT home_id FROM project WHERE project_id = %s", (project_id,))
                project_results = cursor.fetchone()

                home_id = project_results[0]
                cursor.execute("SELECT sector FROM home WHERE home_id = %s", (home_id,))
                home_results = cursor.fetchone()
                sector = home_results[0]
                print('板块', sector)

            except Exception as e:
                if isinstance(e, TypeError) and str(e) == "'NoneType' object is not subscriptable":
                    print('该问题ID下没有数据，跳过')
                    continue
                else:
                    raise ValueError(f'{e}')

            cursor.execute("SELECT material_content, material_order FROM material_new WHERE problem_id = %s", (problem_id,))
            material_results = cursor.fetchall()
            材料内容 = ''
            for ma in material_results:
                材料内容 += ma[0] + '\n'

            cursor.execute("SELECT query_content, query_order FROM query_new WHERE problem_id = %s", (problem_id,))
            query_results = cursor.fetchall()
            提问内容 = ''
            for qu in query_results:
                提问内容 += qu[0] + '\n'

            details_dict = OrderedDict()
            details_dict['问题标题'] = 问题题目.strip()
            details_dict['依据材料的内容'] = 材料内容.strip()
            details_dict['请发行人说明或者披露的内容'] = 提问内容.strip()
            pprint(details_dict)
            test_dict[problem_id]['切片入参'] = details_dict
            test_dict[problem_id]['问题分类'] = 问题分类
            import os

            cursor.close()
        #
        # with open(os.path.join('../data/', 'test.json'), 'w', encoding='utf-8') as f:
        #     json.dump(test_dict, f, ensure_ascii=False, indent=4)
            if sector == '科创板':
                problem_type_dict = innovation_dict
            elif sector == '主板':
                problem_type_dict = motherboard_dict
            elif sector == '创业板':
                problem_type_dict = entrepreneurship_dict
            key_value_list = list(problem_type_dict.items())  # 将字典转换为键值对的列表

            query_message_list = []
            for i in range(0, len(key_value_list), 10):
                batch = key_value_list[i: i + 10]
                type_dict = dict(batch)
                query_message = prompt_problem_classification(details_dict, type_dict)  # 问题字典
                # query_message = prompt_content_classification(content_text, type_dict)  # 招股书片段

                # print(1111111111111111111111111111111111111111)
                # print(query_message)

                query_message_list.append(query_message)

            chat_prompt = ChatPromptTemplate.from_messages(
                [("system", "你是一个证券交易所中审核上市文件的专家。"),
                 ("human", "{query_message}")]
            )
            chain = chat_prompt | self.llm | StrOutputParser()

            problem_type_index_list = []
            for chain_results in await chain.abatch(
                    [{"query_message": query_message} for query_message in query_message_list]):
                print(chain_results)
                try:
                    type_list = eval(Get_List(chain_results))
                except:
                    return f'报错index{problem_id}'
                problem_type_index_list += type_list
            problem_type_index_list = sorted(list(set(problem_type_index_list)))
            if len(problem_type_index_list) == 0:
                print('problem_type_list:', problem_type_index_list)
            else:
                if max(problem_type_index_list) <= 58 and min(problem_type_index_list) >= 0:
                    print('problem_type_list:', problem_type_index_list)
                else:
                    # raise ValueError('问题分类索引不在[0,58]范围内')
                    return f'报错index{problem_id}'

            problem_type_name_list = []
            for type_index in problem_type_index_list:
                problem_type_name_list.append(problem_type_dict[type_index])
            print(problem_type_name_list)
            problem_type_name_list = str(problem_type_name_list)

            DB = LawDB()
            await DB.init_conn()
            DB_text = await DB.update_problem_classification(problem_id, problem_type_name_list)

            e_time = time.time()
            print('单条时间', e_time - s_time)

        db_laws.close()
        return '全部结束'



if __name__ == '__main__':
    import asyncio
    import threading
    Class = Problem_Classification()
    # asyncio.run(Class.classification_problem(3297, 36747))

    results_text = ''
    start_index = 1
    while str(results_text) != '全部结束':
        results_text = asyncio.run(Class.classification_problem(start_index, 36747))
        if str(results_text)[:2] == '报错':
            print(results_text)
            start_index = int(str(results_text).split('报错index')[1])


    # def run_async_function(start_index, end_index):
    #     # start_index = 5615
    #     while True:
    #         try:
    #             # 使用 asyncio.run 来运行异步函数
    #             results_text = asyncio.run(Class.classification_problem(start_index, end_index))
    #             if str(results_text)[:2] == '报错':
    #                 print(results_text)
    #                 start_index = int(str(results_text).split('报错index')[1] + 1)  # 假设报错index后面紧跟的是新的起始索引
    #             else:
    #                 break  # 如果没有报错，退出循环
    #         except Exception as e:
    #             print(f"An error occurred: {e}")
    #             break  # 如
    #
    #     # loop = asyncio.new_event_loop()
    #     # asyncio.set_event_loop(loop)
    #     # loop.run_until_complete(Class.classification_problem(start_index, end_index))
    #     # loop.close()
    #
    #
    # start_index = 36358
    # all_num = start_index + 900 + 1
    # part = 300
    # threads = []
    # for i in range(start_index, all_num, part):
    #     # print(i, i + part)
    #     # part_labor = all_num[i: i + part]
    #     thread = threading.Thread(target=run_async_function, args=(i, i + part,))
    #     print('开始线程：', i)
    #     threads.append(thread)
    #     thread.start()
    #     time.sleep(2)  # 这里的时间间隔可以根据实际情况调整
    #
    # # 等待所有线程完成
    # for thread in threads:
    #     thread.join()
    # print('所有线程完成')


